import tensorflow as tf
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
import datetime

class TaskPredictor:
    def __init__(self):
        self.model = None
        self.scaler = MinMaxScaler()
        self.features = ['day_of_week', 'is_holiday', 'planned_tasks']
        self.target = 'completion_rate'

    def load_data(self, file_path):
        """加载任务数据并进行预处理"""
        df = pd.read_csv(file_path)
        df['date'] = pd.to_datetime(df['date'])
        # 提取特征
        df['day_of_week'] = df['date'].dt.dayofweek + 1  # 1-7表示周一到周日
        # 标准化特征
        df[self.features] = self.scaler.fit_transform(df[self.features])
        return df

    def build_model(self, input_shape):
        """构建神经网络模型"""
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(32, activation='relu', input_shape=input_shape),
            tf.keras.layers.Dense(16, activation='relu'),
            tf.keras.layers.Dense(1, activation='sigmoid')  # 完成率在0-1之间
        ])
        model.compile(optimizer='adam', loss='mse', metrics=['mae'])
        return model

    def train(self, df, epochs=50, batch_size=8):
        """训练模型"""
        X = df[self.features].values
        y = df[self.target].values
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

        self.model = self.build_model((X_train.shape[1],))
        history = self.model.fit(
            X_train, y_train,
            epochs=epochs,
            batch_size=batch_size,
            validation_data=(X_test, y_test),
            verbose=1
        )

        # 评估模型
        loss, mae = self.model.evaluate(X_test, y_test, verbose=0)
        print(f"模型评估结果 - 损失: {loss:.4f}, 平均绝对误差: {mae:.4f}")
        return history

    def predict_completion_rate(self, day_of_week, is_holiday, planned_tasks):
        """预测给定任务量的完成率"""
        if self.model is None:
            raise ValueError("模型尚未训练，请先调用train方法")

        # 准备输入数据
        input_data = np.array([[day_of_week, is_holiday, planned_tasks]])
        input_data_scaled = self.scaler.transform(input_data)

        # 预测完成率
        completion_rate = self.model.predict(input_data_scaled)[0][0]
        return completion_rate

    def recommend_planned_tasks(self, day_of_week, is_holiday, max_tasks=20, target_completion=0.8):
        """推荐最佳任务量以达到目标完成率"""
        best_tasks = 0
        best_rate_diff = float('inf')

        # 搜索最佳任务量
        for tasks in range(1, max_tasks + 1):
            rate = self.predict_completion_rate(day_of_week, is_holiday, tasks)
            rate_diff = abs(rate - target_completion)

            if rate_diff < best_rate_diff:
                best_rate_diff = rate_diff
                best_tasks = tasks

        predicted_rate = self.predict_completion_rate(day_of_week, is_holiday, best_tasks)
        return best_tasks, predicted_rate

if __name__ == "__main__":
    # 示例用法
    predictor = TaskPredictor()
    data = predictor.load_data('../data/task_data_template.csv')
    predictor.train(data)

    # 预测明天（假设明天是周二，非节假日）
    tomorrow = datetime.date.today() + datetime.timedelta(days=1)
    day_of_week = tomorrow.weekday() + 1
    is_holiday = 0  # 需要根据实际节假日表调整

    recommended_tasks, predicted_rate = predictor.recommend_planned_tasks(day_of_week, is_holiday)
    print(f"推荐明天任务量: {recommended_tasks}, 预计完成率: {predicted_rate:.2f}")